import json
import zlib
from datetime import datetime

#
from zoneinfo import ZoneInfo
import lmdb
import cbor2

#
from common import *


LOCAL_TZ = ZoneInfo("Asia/Shanghai")


class CachedbConn:
    # https://lmdb.readthedocs.io/en/release/
    def __init__(self, env_path, map_size, expiration):
        # 打开或创建一个新的LMDB环境
        self._env = lmdb.open(env_path, map_size=map_size, max_dbs=2)
        self._expiration = expiration
        with self._env.begin(write=True) as txn:
            self._env.open_db(key=b"tmdb", txn=txn)
            self._env.open_db(key=b"bgmtv", txn=txn)

    def close(self):
        if self._env is not None:
            self._env.close()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

    def put(self, metaDB_name: str, key: str, obj: bytes):
        with self._env.begin(write=True) as txn:
            db = self._env.open_db(key=metaDB_name.encode(), txn=txn)
            content = {"size": len(obj), "ctime": datetime.now(), "content": obj}
            txn.put(key.encode(), cbor2.dumps(content, timezone=LOCAL_TZ), db=db)

    def put_json(self, metaDB_name: str, key: str, obj: dict):
        obj = zlib.compress(json.dumps(obj, ensure_ascii=False).encode())
        self.put(metaDB_name, key, obj)

    def get(self, metaDB_name: str, key: str):
        with self._env.begin() as txn:
            db = self._env.open_db(key=metaDB_name.encode(), txn=txn)
            data_in_db = txn.get(key.encode(), default=None, db=db)
            if data_in_db is None:
                return None
            return cbor2.loads(data_in_db)["content"]

    def get_json(self, metaDB_name: str, key: str):
        result = self.get(metaDB_name, key)
        return None if result is None else json.loads(zlib.decompress(result).decode())

    def vaccum_by_expiration(self, metaDB_name: str):
        expiration_cond = lambda value: (datetime.now(LOCAL_TZ) - value["ctime"]).days >= self._expiration
        with self._env.begin(write=True) as txn:
            db = self._env.open_db(key=metaDB_name.encode(), txn=txn)
            with txn.cursor(db) as cursor:
                deleting_keys = [key for key, value in cursor.iternext() if expiration_cond(cbor2.loads(value))]
            for key in deleting_keys:
                txn.delete(key, db=db)
